from flask import Flask, request, jsonify
import os
import uuid
import logging
import threading
from werkzeug.utils import secure_filename
from config import Config
from detector import YOLODetector

app = Flask(__name__)
app.config.from_object(Config)

# 设置日志
logging.basicConfig(level=logging.INFO)
app.logger.setLevel(logging.INFO)

# 初始化检测器和锁
detector_current = None
detector_lock = threading.Lock()

def create_detector():
    return YOLODetector(
        model_path=app.config['MODEL_PATH'],
        confidence_thresh=app.config['CONFIDENCE_THRESH']
    )

if detector_current is None:
    app.logger.info("Initializing detector...")
    detector_current = create_detector()

def async_reload_model():
    global detector_current
    try:
        app.logger.info("Reloading model...")
        new_detector = create_detector()
        with detector_lock:
            detector_current = new_detector
        app.logger.info("Model reloaded successfully")
    except Exception as e:
        app.logger.error(f"Model reload failed: {str(e)}", exc_info=True)

@app.route('/reload-model', methods=['POST'])
def reload_model():
    try:
        app.logger.info("Reloading model...")
        thread = threading.Thread(target=async_reload_model)
        thread.start()
        return jsonify({"status": "model reload process started"}), 200
    except Exception as e:
        app.logger.error(f"Model reload failed: {str(e)}", exc_info=True)
        return jsonify({"error": str(e)}), 500

@app.route('/predict', methods=['POST'])
def predict():
    if 'file' not in request.files:
        return jsonify({"error": "No file provided"}), 400

    file = request.files['file']
    if file.filename == '':
        return jsonify({"error": "Empty filename"}), 400

    # 生成唯一文件名避免冲突
    filename = secure_filename(file.filename)
    ext = filename.rsplit('.', 1)[-1].lower()
    if ext not in app.config['ALLOWED_EXTENSIONS']:
        return jsonify({"error": "Unsupported file format"}), 400

    unique_filename = f"{uuid.uuid4().hex}_{filename}"
    file_path = os.path.join(app.config['UPLOAD_FOLDER'], unique_filename)

    try:
        # 保存上传文件
        file.save(file_path)

        # 获取并验证参数，必须存在且有效
        try:
            flight_height = float(request.form['flight_height'])
            sensor_width = float(request.form['sensor_width'])
            sensor_height = float(request.form['sensor_height'])
            focal_length = float(request.form['focal_length'])
        except (KeyError, ValueError):
            return jsonify({"error": "Missing or invalid sensor parameters"}), 400

        if flight_height <= 0 or sensor_width <= 0 or sensor_height <= 0 or focal_length <= 0:
            return jsonify({
                "error": "Flight height, sensor dimensions and focal length must be positive values"
            }), 400

        # 位置参数可选，默认0
        longitude = float(request.form.get('longitude', 0))
        latitude = float(request.form.get('latitude', 0))
        altitude = float(request.form.get('altitude', 0))
        pitch_angle = float(request.form.get('pitch_angle', 0))  # 俯仰角（度）
        roll_angle = float(request.form.get('roll_angle', 0))  # 横滚角（度）

        # 使用当前模型进行推理，线程安全读取
        with detector_lock:
            current_detector = detector_current

        # 根据文件类型处理
        if ext in ('png', 'jpg', 'jpeg'):
            result = current_detector.predict_image(
                file_path,
                flight_height,
                sensor_width,
                sensor_height,
                focal_length,
                pitch_angle,
                roll_angle
            )
            result["type"] = "image"

        elif ext in ('mp4', 'avi', 'mov'):
            result = current_detector.predict_video(
                file_path,
                flight_height,
                sensor_width,
                sensor_height,
                focal_length,
                pitch_angle,
                roll_angle
            )
            result["type"] = "video"

        else:
            return jsonify({"error": "Unsupported file format"}), 400

        # 添加模型信息
        result["model"] = {
            "name": os.path.basename(app.config['MODEL_PATH']),
            "version": app.config.get('MODEL_VERSION', '1.0'),
            "conf_threshold": float(app.config['CONFIDENCE_THRESH'])
        }

        # 添加位置信息
        result["Original location"] = {
            "longitude": longitude,
            "latitude": latitude,
            "altitude": altitude
        }

        # 添加姿态信息
        result["orientation"] = {
            "pitch_angle": pitch_angle,
            "roll_angle": roll_angle
        }

        # 增加图片名称字段
        result["image_name"] = filename

        # 计算并添加修正后的经纬度
        if isinstance(result.get('detections', []), list) and len(result['detections']) > 0:
            # 获取图像尺寸
            image_width, image_height = result.get('original_size', (0, 0))
            gsd_x, gsd_y = result.get('gsd_x', 0), result.get('gsd_y', 0)

            if image_width == 0 or image_height == 0:
                app.logger.warning("Image size is zero, skipping coordinate correction")
            else:
                # 为每个检测结果计算修正后的经纬度
                for detection in result['detections']:
                    # 获取检测框中心点
                    bbox = detection['bbox']
                    center_x = (bbox['xmin'] + bbox['xmax']) / (2 * image_width)
                    center_y = (bbox['ymin'] + bbox['ymax']) / (2 * image_height)

                    # 计算修正后的经纬度
                    corrected_coords = current_detector.calculate_corrected_coordinates(
                        center_x, center_y, image_width, image_height, gsd_x, gsd_y, longitude, latitude
                    )

                    # 添加修正后的经纬度到检测结果
                    detection['location'] = corrected_coords

        # 检查检测结果是否为空
        if isinstance(result.get('detections', []), list) and len(result['detections']) == 0:
            return jsonify(result), 205

        return jsonify(result), 200

    except Exception as e:
        app.logger.error(f"Prediction failed: {str(e)}", exc_info=True)
        return jsonify({"error": "Internal server error"}), 500

    finally:
        # 清理临时文件
        if os.path.exists(file_path):
            try:
                os.remove(file_path)
                app.logger.info(f"Removed temporary file: {file_path}")
            except Exception as e:
                app.logger.error(f"Error removing file {file_path}: {str(e)}")

@app.route('/update-model', methods=['POST'])
def update_model():
    if 'file' not in request.files:
        return jsonify({"error": "No model file provided"}), 400

    file = request.files['file']
    if file.filename == '':
        return jsonify({"error": "Empty filename"}), 400

    filename = secure_filename(file.filename)
    ext = filename.rsplit('.', 1)[-1].lower()
    allowed_model_exts = {'pt', 'pth', 'onnx', 'yaml'}
    if ext not in allowed_model_exts:
        return jsonify({"error": "Unsupported model file format"}), 400

    unique_filename = f"{uuid.uuid4().hex}_{filename}"
    new_model_path = os.path.join(app.config['UPLOAD_FOLDER'], unique_filename)

    try:
        # 保存新的模型文件
        file.save(new_model_path)

        # 更新配置文件中的模型路径，线程安全
        with detector_lock:
            app.config['MODEL_PATH'] = new_model_path

        # 异步重新加载模型
        thread = threading.Thread(target=async_reload_model)
        thread.start()

        return jsonify({"status": "model update process started"}), 200

    except Exception as e:
        app.logger.error(f"Model update failed: {str(e)}", exc_info=True)
        return jsonify({"error": str(e)}), 500

# 确保上传文件夹存在（无论何种启动方式）
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)

if __name__ == '__main__':
    app.logger.info(f"Starting server on {app.config['HOST']}:{app.config['PORT']}")
    app.logger.info(f"Upload folder: {app.config['UPLOAD_FOLDER']}")
    app.logger.info(f"Using model: {app.config['MODEL_PATH']}")
    app.run(host=app.config['HOST'], port=app.config['PORT'], debug=app.config['DEBUG'])